import asyncio
from typing import Dict, List, Optional

from llm import LLM
from prompts.summary_prompt import CHUNK_SUMMARY_PROMPT, PAPER_SUMMARY_PROMPT
from query import Query


class SummaryExtractor:
    """A class to extract and generate summaries from paper chunks using LLM.

    Attributes:
        llm (LLM): Language model instance for generating summaries
        query (Query): Query instance for retrieving paper chunks
    """

    def __init__(self, llm: LLM):
        self.llm = llm
        self.query = Query()

    async def summarize_paper(
        self, paper_id: Optional[str], paper_title: Optional[str]
    ) -> str:
        """Generate a complete paper summary from its chunks.

        Args:
            paper_id (str): Unique identifier of the paper
            paper_title (str): Title of the paper

        Returns:
            str: Generated summary of the entire paper

        Raises:
            ValueError: If both paper_id and paper_title are empty
            Exception: If chunk retrieval or summary generation fails
        """
        try:
            if not paper_id and not paper_title:
                raise ValueError("Either paper_id or paper_title must be provided")

            chunks = (
                self.query.query_by_id(paper_id)
                if paper_id
                else self.query.query_by_title(paper_title)
            )
            if not chunks:
                raise ValueError(
                    f"No chunks found for paper_id={paper_id} or title={paper_title}"
                )

            chunk_summaries = await self.summarize_chunks(chunks)
            prompt = PAPER_SUMMARY_PROMPT.format(chunk_summaries=chunk_summaries)
            return await self.llm.completion(prompt)
        except Exception as e:
            raise Exception(f"Failed to summarize paper: {str(e)}")

    async def summarize_chunks(self, chunks: List[Dict]) -> List[str]:
        """Summarize multiple chunks concurrently.

        Args:
            chunks (List[Dict]): List of paper chunks to summarize

        Returns:
            List[str]: List of generated summaries for each chunk

        Raises:
            Exception: If chunk summarization fails
        """
        try:
            tasks = [self.summarize_chunk(chunk) for chunk in chunks]
            return await asyncio.gather(*tasks)
        except Exception as e:
            raise Exception(f"Failed to summarize chunks: {str(e)}")

    async def summarize_chunk(self, chunk: Dict) -> str:
        """Summarize a single chunk of text.

        Args:
            chunk (Dict): Dictionary containing chunk text and metadata

        Returns:
            str: Generated summary for the chunk

        Raises:
            Exception: If chunk summarization fails
        """
        try:
            prompt = CHUNK_SUMMARY_PROMPT.format(chunk=chunk)
            return await self.llm.completion(prompt)
        except Exception as e:
            raise Exception(f"Failed to summarize chunk: {str(e)}")


if __name__ == "__main__":
    llm = LLM()
    summary_extractor = SummaryExtractor(llm)
    summary = asyncio.run(
        summary_extractor.summarize_paper("651b7dbc3fda6d7f06304579", None)
    )
    print(summary)
